#include <mpi.h>
#include <omp.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/timeb.h>
#include <time.h>
#include "dictionary.h"
#include "rdictionary.h"

#define ROOT_PROCESS 0
#define RECV_TIMEOUT 500
#define MAX_WORD_SIZE 5

/* Displays the message and abort the program. */
void error(const char *message);

/* Returns the number of milliseconds since Epoch. */
long current_time_millis();

/* Returns the number of seconds since 'start'. */
float current_time_seconds(long start);

/* Randomly generates a letter. */
char random_letter();

/* Randomly generates a word. */
void generate_word(char *word);

/* Copies the file contents in a buffer. */
char *read_file(char *filename);

/* Core function to find words and communicate. */
void find_words(int rank, int num_processes, char *filename);

/* Core function executed by root that receives the found words. */
void count_words(char *filename);

int main(int argc, char *argv[])
{
    int rank, num_processes;

    /* Checks for file argument. */
    if (argc < 2) {
        fprintf(stderr, "Missing input file.\n");
        return 1;
    }

    if (MPI_Init(&argc, &argv)) {
        fprintf(stderr, "MPI: initialization error.\n");
        return 1;
    }

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &num_processes);

    if (rank == ROOT_PROCESS)
        /* Root process. */
        count_words(argv[1]);
    else
        /* Other processes. */
        find_words(rank, num_processes, argv[1]);

    MPI_Finalize();

    return 0;
}

void error(const char *message)
{
    fprintf(stderr, "%s\n", message);
    MPI_Abort(MPI_COMM_WORLD, 1);
}

long current_time_millis()
{
    struct timeb tp;

    ftime(&tp);
    return tp.time * 1000 + tp.millitm;
}

float current_time_seconds(long start)
{
    return (current_time_millis() - start) / 1000.0f;
}

char random_letter()
{
    return 'a' + rand() % 26;
}

void generate_word(char *word)
{
    int i;

    /*
     * Generates random letters with an increasing chance of stopping before the
     * end.
     */
    for (i = 0; i < MAX_WORD_SIZE; ++i) {
        if ((double) rand() / RAND_MAX < 0.01 * i)
            break;
        word[i] = random_letter();
    }

    /* Fill the rest of the generated word with underscores. */
    for (; i < MAX_WORD_SIZE; ++i)
        word[i] = '_';
    word[i] = '\0';
}

char *read_file(char *filename)
{
    int ret;
    char *buffer;

    MPI_Offset file_size;
    MPI_File file;
    MPI_Status status;

    ret = MPI_File_open(MPI_COMM_WORLD, filename, MPI_MODE_RDONLY, MPI_INFO_NULL, &file);
    if (ret)
        error("MPI: file open error.");

    MPI_File_get_size(file, &file_size);
    buffer = (char *) malloc(sizeof(char) * file_size);
    if (buffer == NULL)
        error("Out of memory.");

    ret = MPI_File_read_all(file, buffer, (int) file_size, MPI_CHAR, &status);
    if (ret)
        error("MPI: file read error");

    MPI_File_close(&file);

    return buffer;
}

void find_words(int rank, int num_processes, char *filename)
{
    int word_index, tid;
    char word[MAX_WORD_SIZE + 1];
    char *input_buffer, *word_ptr;
    int start_millis, current_millis;

    /* Reads file. */
    input_buffer = read_file(filename);

    /*
     * Initializes the dictionary of words.
     * The dictionary will manage all parts of all words.
     */
    dict_init();

    /* Adds each word in the dictionary. */
    word_ptr = strtok(input_buffer, ",");
    while (word_ptr != NULL) {
        dict_add(word_ptr);
        word_ptr = strtok(NULL, ",");
    }

    free(input_buffer);

    /* Optimizes the dictionary in order to speed up mark operations. */
    dict_optimize();

    MPI_Status status;
    MPI_Request request;
    int done;

    #pragma omp parallel \
        private(word_index, word, tid, status, request, done)
    {
        tid = omp_get_thread_num();

        /* The first thread will receive parts of words found by other processes. */
        if (tid == 0) {
            while (!dict_is_empty()) {
                #pragma omp critical (mpi)
                MPI_Irecv(&word_index, 1, MPI_INT, MPI_ANY_SOURCE, 0, MPI_COMM_WORLD, &request);

                /*
                 * Waits until operation is done or
                 * cancels operation if timeout is reached.
                 */
                start_millis = current_time_millis();
                do {
                    #pragma omp critical (mpi)
                    MPI_Test(&request, &done, &status);
                    current_millis = current_time_millis();
                } while (!done && current_millis - start_millis < RECV_TIMEOUT);

                if (!done && current_millis - start_millis >= RECV_TIMEOUT) {
                    MPI_Cancel(&request);
                    continue;
                }

                /* Marks the found part of words using its index. */
                #pragma omp critical (mark)
                dict_mark(word_index);
            }
        /* Other threads will look for words. */
        } else {
            srand(rank * 1000 + tid);
            while (!dict_is_empty()) {
                int i;

                /*
                 * Generates a random word with at most 5 letters.
                 * Verifies if it exists in the dictionary.
                 */
                generate_word(word);
                word_index = dict_search(word);
                if (word_index == -1)
                    continue;

                /* If the part of word is new, marks and sends it. */
                for (i = 0; i < num_processes; ++i)
                    if (i != rank) {
                        #pragma omp critical (mpi)
                        MPI_Isend(&word_index, 1, MPI_INT, i, 0, MPI_COMM_WORLD, &request);

                        do {
                            #pragma omp critical (mpi)
                            MPI_Test(&request, &done, &status);
                        } while (!done);
                    } else {
                        #pragma omp critical (mark)
                        dict_mark(word_index);
                    }
            }
        }
    }

    /* Deletes the dictionary. */
    dict_delete();
}

void count_words(char *filename)
{
    int word_index, percent, objective;
    char *input_buffer, *word_ptr;
    MPI_Status status;
    long start_millis;

    start_millis = current_time_millis();
    printf("Reading file...\n");

    /* Reads file. */
    input_buffer = read_file(filename);

    printf("Done (%fs).\nBuilding tree...\n", current_time_seconds(start_millis));

    /*
     * Initializes the root dictionary.
     * The root dictionary knows which word each part belongs.
     */
    rdict_init(62000);

    /* Adds each word in the root dictionary. */
    word_ptr = strtok(input_buffer, ",");
    while (word_ptr != NULL) {
        rdict_add(word_ptr);
        word_ptr = strtok(NULL, ",");
    }

    free(input_buffer);

    printf("Tree built (%fs).\nOptimizing tree for retrieval...\n", current_time_seconds(start_millis));

    /* Optimizes the dictionary in order to speed up mark operations. */
    rdict_optimize();

    printf("Done (%fs).\nSearching words...\n", current_time_seconds(start_millis));

    objective = 10;
    percent = 0;

    /* Receives parts of words found by other processes. */
    while (percent < 100) {
        MPI_Recv(&word_index, 1, MPI_INT, MPI_ANY_SOURCE, 0, MPI_COMM_WORLD, &status);

        /* Checks if the found part complete some word or words. */
        if (rdict_mark(word_index) == 0)
            continue;

        /* Calculates and displays the percentage of found words. */
        percent = (int) (((float) rdict_mark_count / rdict_word_count) * 100);
        if (percent >= objective) {
            printf("Reached %d%%: (%fs)\n", objective, current_time_seconds(start_millis));
            objective += 10;
        }
    }

    /* Deletes the root dictionary. */
    rdict_delete();
}
